import argparse
import os
import time
from datetime import timedelta

from log_manager import LogManager


class OriginalVideo(object):
    pass


def split_text_by_max_length(input_file, max_length):
    with open(input_file, "r", encoding="utf-8") as f:
        text = f.read()
    paragraphs = text.split("\n")
    split_text = []
    current_length = 0  # Track the length of the accumulated text.
    for paragraph in paragraphs:
        paragraph = paragraph.strip()
        if not paragraph:
            continue
        para_length = len(paragraph)
        if current_length + para_length + (1 if split_text else 0) <= max_length:
            # Add 1 to account for the newline character when joining.
            if "原视频" in paragraph:
                yield "\n".join(split_text)
                yield OriginalVideo()
                split_text = []
                current_length = 0
            else:
                split_text.append(paragraph)
                current_length += para_length + (1 if split_text else 0)
        else:
            yield "\n".join(split_text)
            split_text = [paragraph]
            current_length = para_length
    if split_text:  # Yield any remaining text.
        yield "\n".join(split_text)


def format_time(td):
    """将timedelta对象转换为SRT时间格式"""
    total_seconds = td.total_seconds()
    hours = int(total_seconds // 3600)
    minutes = int((total_seconds % 3600) // 60)
    seconds = int(total_seconds % 60)
    milliseconds = int((total_seconds * 1000) % 1000)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"


def main():
    parser = argparse.ArgumentParser(description="Split text file into multiple files")
    parser.add_argument("input_file", help="Input file path")
    parser.add_argument(
        "-m",
        "--max_length",
        type=int,
        default=500,
        help="Max length of each split file",
    )
    args = parser.parse_args()

    input_file = args.input_file
    max_length = args.max_length

    filename, _ = os.path.splitext(input_file)
    current_time = int(time.time())
    output_file = f"{filename}_{current_time}.srt"
    current_time = timedelta(hours=0, minutes=0, seconds=0, milliseconds=0)
    with open(output_file, "w", encoding="utf-8") as f:
        index = 0
        for split in split_text_by_max_length(input_file, max_length):
            if isinstance(split, OriginalVideo):
                current_time += timedelta(seconds=20)
                continue
            index += 1
            text_length = len(split)
            time_length = text_length * 200
            start_time_str = format_time(current_time)
            current_time += timedelta(milliseconds=time_length)
            end_time_str = format_time(current_time)
            context = f"{index}\n{start_time_str} --> {end_time_str}\n"
            f.write(context)
            f.write(split)
            f.write("\n\n")
        LogManager.log(f"Saved to {output_file}")


if __name__ == "__main__":
    main()
